
import sys
import time
import md5
from feature import createFeatureParser
from featureindex import FeatureIndex

FILENAME_KEY = "FILENAME"
TIMESTAMP_KEY = "TIMESTAMP"
INTERVAL_KEY = "INTERVAL"
COUNT_KEY = "FEATURE_COUNT"
KEYS = [FILENAME_KEY, TIMESTAMP_KEY, INTERVAL_KEY, COUNT_KEY]

class BisectionIndex(FeatureIndex):
  '''
  Represents an index for a file containing features sorted by their start position.
  Provides methods to validate an index against its source file, write an index
  to a file, read an index from a file, and retrieve all the features that overlap
  a particular location.
  '''
  def getFeatures(self, query=(None, None, None)):
    '''
    Return an iterator over features overlapping the specified
    query interval, given in the form (chr, startPos, endPos).
    Any of chr, startPos, and endPos may be given as None to
    drop the corresponding restriction; e.g., (None, None, None)
    returns all features.
    '''
    (qchr, qmin, qmax) = query
    if qchr == None:
      # if None, iterate over all chromosomes in index
      for chr in self.chrList:
        for feature in self.getFeatures((chr, qmin, qmax)):
          yield feature
      return
    featureFilename = self.ixByChr[FILENAME_KEY]
    featureParser = createFeatureParser(featureFilename)
    maxFeatureSize = featureParser.getMaxFeatureSize()
    
    # print >> sys.stderr, "# Performing query", query
    ix = self.ixByChr.get(qchr, [])
    f = open(featureFilename, "r")

    minFilePos = long(self.ixByChr[qchr][0].strip().split("\t")[2]) - 1
    maxFilePos = long(self.ixByChr[self.chrList[self.chrList.index(qchr) + 1]][0].strip().split("\t")[2]) - 1

    # Assert that minFilePos is early enough to read a line, then start reading features,
    # i.e., need feature start < qmin - maxFeatureSize or at the start of the chromosome.
    # Otherwise push minFilePos as far right as possible.
    while maxFilePos > minFilePos + 1000:
      midPos = (minFilePos + maxFilePos) / 2
      f.seek(midPos)
      f.readline()
      feature = None
      while not feature:
        feature = featureParser.parse(f.readline())
      (chr, minPos, maxPos) = (feature.chr, feature.spos, feature.epos)
      if chr!=qchr:
        if self.chrList.index(chr) > self.chrList.index(qchr):
          maxFilePos = midPos
        else:
          minFilePos = midPos
      else:
        if (minPos < qmin - maxFeatureSize):
          minFilePos = midPos
        else:
          maxFilePos = midPos

    f.seek(minFilePos)
    f.readline()
    while True:
      line = f.readline()
      if line=="": break
      f.flush()
      feature = featureParser.parse(line)
      if feature==None: continue
      (chr, minPos, maxPos) = (feature.chr, feature.spos, feature.epos)
      if chr!=qchr or (qmax != None and minPos > qmax): break
      if (qmin != None and maxPos < qmin): continue
      yield feature
    f.close()

def validateIndex(filename):
  index = FeatureIndex(filename)
  err = index.validationError()
  if err != None:
    raise Exception(err)
  
def getFeatures(filename, query=(None, None, None)):
  '''Read features from an indexed file, falling back to unindexed file if needed.'''
  if not filename.endswith("INDEX"):
    print >> sys.stderr, filename, "does not end with INDEX; attempting (slow) read of unindexed features..."
    features = _getUnindexedFeatures(filename, query)
  else:
    index = FeatureIndex(filename)
    features = index.getFeatures(query)
  for feature in features:
    yield feature
  return

def _getUnindexedFeatures(filename, query=(None, None, None)):
  '''
  Fallback method for reading features from original, unindexed file.
  Query given as (chr, startPos, endPos); if any of these are None,
  the corresponding restriction is droppped.
  '''
  featureParser = createFeatureParser(filename)
  (qchr, qmin, qmax) = query
  f = open(filename, "r")
  foundBlock = False
  for line in f:
    feature = featureParser.parse(line)
    if feature==None: continue
    (chr, minPos, maxPos) = (feature.chr, feature.spos, feature.epos)
    if (qchr!=None and chr!=qchr) or (qmax != None and minPos > qmax):
      if foundBlock: break
      continue
    if (qmin != None and maxPos < qmin): continue
    foundBlock = True
    yield feature
  f.close()

if __name__=="__main__":
  if sys.argv[1]=="-I":
    filename = sys.argv[2]
    index = FeatureIndex()
    index.createFromFeatures(filename)
    index.printIndex()
  else:
    indexFilename = sys.argv[1]
    if len(sys.argv)<4:
      validateIndex(indexFilename)
    else:
      minPos = maxPos = long(sys.argv[3])
      if len(sys.argv)>4: maxPos = long(sys.argv[4])
      query = (sys.argv[2], minPos, maxPos)
      features = getFeatures(indexFilename, query)
      for f in features:
        print (f.chr, f.spos, f.epos, f.content.strip())

